from preprocess.openpose.run_openpose import OpenPose
from PIL import Image, ImageDraw
import numpy as np
import os

openpose_ckp = '/mnt/nas/shengjie/huggingface_model_local/ootd/checkpoints/openpose/ckpts'
pose_shape = (384, 512)

def get_pose_predictor():
    # 检查模型路径是否存在
    if not os.path.exists(openpose_ckp):
        raise FileNotFoundError(f"OpenPose checkpoint path does not exist: {openpose_ckp}")
    openpose_model = OpenPose(0, openpose_ckp)
    return openpose_model

def get_pose_res(pose_predictor=None,img_path='',img_pil=None,pose_shape=(384,512)):
    img = img_pil if img_pil is not None else Image.open(img_path)
    keypoints = pose_predictor( img.resize( pose_shape ) )

    # print(keypoints)
    return keypoints

def get_pose_pil_by_keypoints(pose_data, target_shape):
    # 配置参数
    width, height = target_shape
    line_width = 3
    point_radius = 5
    bg_color = (255, 255, 255)  # 白色背景
    skeleton_color = (0, 0, 255)  # 蓝色骨架
    joint_color = (255, 0, 0)  # 红色关节点

    # 骨架连接关系 (COCO格式)
    POSE_PAIRS = [
        (1, 0), (1, 2), (1, 5),  # 躯干和肩膀
        (2, 3), (3, 4),          # 右臂
        (5, 6), (6, 7),          # 左臂
        (1, 8), (1, 11),         # 髋部连接
        (8, 9), (9, 10),         # 右腿
        (11, 12), (12, 13)       # 左腿
    ]

    # 创建图像
    pose_img = Image.new('RGB', (width, height), bg_color)
    draw = ImageDraw.Draw(pose_img)

    # 获取关键点坐标
    if not isinstance(pose_data, dict) or "pose_keypoints_2d" not in pose_data:
        raise ValueError("pose_data must be a dict with 'pose_keypoints_2d' key.")
    pose_array = np.array(pose_data["pose_keypoints_2d"])
    if pose_array.ndim == 1:
        if pose_array.size % 2 != 0:
            raise ValueError("pose_keypoints_2d must have even length.")
        pose_array = pose_array.reshape((-1, 2))
    elif pose_array.ndim == 2 and pose_array.shape[1] == 2:
        pass
    else:
        raise ValueError("pose_keypoints_2d must be shape (N,2) or flat list of even length.")

    # 绘制骨架连线
    for pair in POSE_PAIRS:
        start_idx, end_idx = pair
        if start_idx >= len(pose_array) or end_idx >= len(pose_array):
            continue
        start_point = tuple(pose_array[start_idx])
        end_point = tuple(pose_array[end_idx])
        # 过滤无效点 (坐标为0或负数)
        if all(coord > 0 for coord in start_point) and all(coord > 0 for coord in end_point):
            draw.line([start_point, end_point], fill=skeleton_color, width=line_width)

    # 绘制关节点
    for point in pose_array:
        x, y = point
        if x > 0 and y > 0:  # 过滤无效点
            draw.ellipse(
                [(x - point_radius, y - point_radius),
                 (x + point_radius, y + point_radius)],
                fill=joint_color
            )

    return pose_img

if __name__ == '__main__':
    predictor = get_pose_predictor()
    img_path = './tmp_tryon.jpg'
    if not os.path.isfile(img_path):
        raise FileNotFoundError(f"Image file not found: {img_path}")
    keypoints = get_pose_res(predictor, img_path=img_path, pose_shape=pose_shape)
    pose_img = get_pose_pil_by_keypoints(keypoints, pose_shape)
    pose_img.save('tmp_tryon2.jpg')